package cn.com.jit;

import java.util.ArrayList;
import java.util.List;
import java.util.regex.Matcher;
import java.util.regex.Pattern;

public class SQLColumnExtractor {

    /**
     * 从 SELECT 语句中提取列名，并去掉表别名。
     *
     * @param sql SQL 语句
     * @return 列名列表
     */
    public static List<String> extractColumnNames(String sql) {
        List<String> columnNames = new ArrayList<>();

        // 正则表达式匹配 SELECT 语句中的列名
        String regex = "SELECT\\s+(?:DISTINCT\\s+)?(?:\\*|((?:[a-zA-Z0-9_`\"'.]+(?:\\.[a-zA-Z0-9_`\"'.]+)?\\s*(?:AS\\s+[a-zA-Z0-9_`\"'.]+)?(?:,\\s*)?)+))\\s+FROM";

        Pattern pattern = Pattern.compile(regex, Pattern.CASE_INSENSITIVE);
        Matcher matcher = pattern.matcher(sql);

        while (matcher.find()) {
            String columnsPart = matcher.group(1);
            if (columnsPart != null) {
                // 分割列名字符串
                String[] columns = columnsPart.split(",\\s*");
                for (String column : columns) {
                    // 去除多余的空格
                    column = column.trim();
                    // 去除引号
                    column = column.replaceAll("^['\"]|['\"]$", "");

                    // 去掉表别名
                    int dotIndex = column.indexOf('.');
                    if (dotIndex != -1) {
                        column = column.substring(dotIndex + 1);
                    }

                    // 去掉 AS 后的别名
                    int asIndex = column.indexOf(" AS ");
                    if (asIndex != -1) {
                        column = column.substring(0, asIndex).trim();
                    }

                    columnNames.add(column);
                }
            }
        }

        return columnNames;
    }

    public static void main(String[] args) {
        String sql = "Select id, name, address FROM users WHERE id > 10";
        List<String> columnNames = extractColumnNames(sql);
        System.out.println("Extracted Column Names: " + columnNames);

        String sqlWithTable = "SELECT users.id AS user_id, users.name AS user_name, addresses.address AS addr FROM users INNER JOIN addresses ON users.id = addresses.user_id";
        List<String> columnNamesWithTable = extractColumnNames(sqlWithTable);
        System.out.println("Extracted Column Names with Table: " + columnNamesWithTable);

        String sqlWithQuotes = "SELECT \"id\", 'name', `address` FROM users";
        List<String> columnNamesWithQuotes = extractColumnNames(sqlWithQuotes);
        System.out.println("Extracted Column Names with Quotes: " + columnNamesWithQuotes);
    }
}
